#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
转换旧格式的HDF5文件为新的整合格式

旧格式：每个group_i单独存储
新格式：所有数据合并为大矩阵存储
"""

import os
import numpy as np
import h5py
from utils.h5data import read_h5_data, save_h5_data
from utils.h5data import MemoryMappedH5
import argparse

def convert_file(old_file, new_file):
    """
    将旧格式的HDF5文件转换为新格式
    
    参数:
    old_file: 旧格式HDF5文件路径
    new_file: 新格式H5文件路径
    """
    print(f"正在转换文件: {old_file}")
    
    # 使用MemoryMappedH5读取文件以提高性能
    file_io =  MemoryMappedH5(old_file)
    h5_file = file_io.h5_file
    
    # 读取元数据
    metadata, attrs = read_h5_data(keys=['OUTFILE', 'INFILE', 'MARKER', 'EMARKER', 'DATE', 'BASE_STAT'],
                                    group_name='metadata', 
                                    h5file=h5_file,
                                    read_attrs=True)
    
    # 读取统计信息和组信息
    stats_data = read_h5_data(keys=['stats'], h5file=h5_file)[0]
    all_groups = read_h5_data(keys=['all_groups'], h5file=h5_file)[0]
    vinfo = read_h5_data(keys=['vinfo'], h5file=h5_file)[0]
    
    # 解码字符串数据
    all_groups = [g.decode() if isinstance(g, bytes) else g for g in all_groups]
    stats_data = [s.decode() if isinstance(s, bytes) else s for s in stats_data]
    
    print(f"总共找到 {len(all_groups)} 个组")
    
    # 初始化列表用于收集所有数据
    corr_list = []
    t_list = []
    R_list = []
    S_list = []
    te_list = []
    
    # 遍历所有组，收集数据
    for i, group_name in enumerate(all_groups):
        if i % 100 == 0:
            print(f"处理进度: {i}/{len(all_groups)}")
            
        try:
            # 读取组数据
            group_data = read_h5_data(keys=['corr', 't', 'R', 'S', 'te'], 
                                        group_name=group_name,
                                        h5file=h5_file)
            
            corr, t, R, S, te = group_data
            
            # 添加到列表
            corr_list.append(corr)
            t_list.append(t)
            R_list.append(R)
            S_list.append(S)
            te_list.append(te)
            
        except Exception as e:
            print(f"处理组 {group_name} 时出错: {e}")
            continue
    
    # 将数据合并为大矩阵
    if corr_list:
        print("正在合并数据...")
        all_corr = np.stack(corr_list)
        all_t = np.stack(t_list)
        all_R = np.stack(R_list)
        all_S = np.array(S_list)
        all_te = np.array(te_list)
        
        # 创建新文件并保存所有数据
        print("正在保存到新文件...")
        data_dict = {
            'corr': all_corr,
            't': all_t,
            'R': all_R,
            'S': all_S,
            'te': all_te,
            'all_groups': np.array(all_groups, dtype='S'),
            'vinfo': vinfo,
            'stats': np.array(stats_data, dtype='S')
        }
        
        # 保存到新文件
        save_h5_data(new_file, data_in_dict=data_dict, mode='w')
        
        # 复制元数据
        save_h5_data(new_file, 
                    data_in_dict={
                        'OUTFILE': metadata[0],
                        'INFILE': metadata[1],
                        'MARKER': metadata[2],
                        'EMARKER': metadata[3],
                        'DATE': metadata[4],
                        'BASE_STAT': metadata[5]
                    },
                    group_name='metadata',
                    attrs_dict=attrs,
                    mode='a')
        
        print(f"转换完成: {new_file}")
        print(f"数据形状 - corr: {all_corr.shape}, t: {all_t.shape}")
    else:
        print("没有数据需要转换")
        # 即使没有数据也要保存基本的元数据
        save_h5_data(new_file, 
                    data_in_dict={
                        'OUTFILE': metadata[0],
                        'INFILE': metadata[1],
                        'MARKER': metadata[2],
                        'EMARKER': metadata[3],
                        'DATE': metadata[4],
                        'BASE_STAT': metadata[5],
                        'all_groups': np.array(all_groups, dtype='S'),
                        'vinfo': vinfo,
                        'stats': np.array(stats_data, dtype='S')
                    },
                    group_name='metadata',
                    attrs_dict=attrs,
                    mode='w')
        print(f"元数据保存完成: {new_file}")
    file_io.close()

def batch_convert(input_dir, output_dir, pattern=".h5"):
    """
    批量转换目录中的所有HDF5文件
    
    参数:
    input_dir: 输入目录
    output_dir: 输出目录
    pattern: 文件匹配模式
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 查找所有匹配的文件
    files = [f for f in os.listdir(input_dir) if f.endswith(pattern)]
    
    print(f"找到 {len(files)} 个文件需要转换")
    
    for filename in files:
        old_file = os.path.join(input_dir, filename)
        new_file = os.path.join(output_dir, filename)
        
        try:
            convert_file(old_file, new_file)
        except Exception as e:
            print(f"转换文件 {filename} 时出错: {e}")

def main():
    parser = argparse.ArgumentParser(description='将旧格式的HDF5文件转换为新格式')
    parser.add_argument('-i', '--input', required=True, help='输入文件或目录')
    parser.add_argument('-o', '--output', required=True, help='输出文件或目录')
    parser.add_argument('-b', '--batch', action='store_true', help='批量转换目录中的所有文件')
    parser.add_argument('-p', '--pattern', default='.h5', help='文件匹配模式 (默认: .h5)')
    
    args = parser.parse_args()
    
    if args.batch:
        batch_convert(args.input, args.output, args.pattern)
    else:
        convert_file(args.input, args.output)

if __name__ == '__main__':
    main()